from PySide6.QtWidgets import (QApplication, QMainWindow, QTableWidget,
                             QTableWidgetItem, QPushButton, QVBoxLayout, 
                             QHBoxLayout, QWidget, QHeaderView, QMessageBox)
from PySide6.QtCore import Qt, QEvent
from PySide6.QtGui import QColor, QUndoStack, QUndoCommand, QKeySequence, QAction, QPen, QPainter

from collections import defaultdict
from utils.repeat_sequence import mark_repeating_sequences

import pandas as pd
import numpy as np
import random

class MarkCycleCommand(QUndoCommand):
    def __init__(self, table, start_row, end_row, cycle_num, color, prev_marks, prev_colors):
        super().__init__()
        self.table = table
        self.start_row = start_row
        self.end_row = end_row
        self.cycle_num = cycle_num
        self.color = color  # 存储当前循环的颜色
        self.prev_marks = prev_marks  # 存储操作前的标记状态
        self.prev_colors = prev_colors  # 存储操作前的颜色状态

    def redo(self):
        for row in range(self.start_row, self.end_row + 1):
            item = QTableWidgetItem(f"小循环{self.cycle_num}")
            item.setBackground(self.color)
            self.table.setItem(row, 2, item)

    def undo(self):
        for row in range(self.start_row, self.end_row + 1):
            if row in self.prev_marks:
                item = QTableWidgetItem(self.prev_marks[row])
                item.setBackground(self.prev_colors.get(row, QColor(255, 255, 255)))
                self.table.setItem(row, 2, item)
            else:
                self.table.setItem(row, 2, QTableWidgetItem(""))

class WorkstepCycleTool(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("工步序列循环划分工具")
        self.undo_stack = QUndoStack()
        self.cycle_colors = {}  # 存储每个循环的颜色
        self.show_borders = False  # 控制边框是否显示的标志
        self.special_row_ranges = []
        # 创建一个柔和的颜色池，避免过于刺眼的颜色
        self.color_pool = [
            QColor(144, 238, 144),  # 浅绿
            QColor(173, 216, 230),  # 浅蓝
            QColor(255, 228, 196),  # 浅橙
            QColor(230, 230, 250),  # 淡紫
            QColor(255, 255, 224),  # 淡黄
            QColor(255, 182, 193),  # 浅粉
            QColor(188, 143, 143),  # 玫瑰木色
            QColor(204, 255, 255),  # 浅青
            QColor(255, 240, 245),  # 淡粉红
            QColor(240, 255, 240)   # 淡绿
        ]

        # 创建表格
        self.table = QTableWidget(0, 4)  # 3列：工步序号、工步状态、标记
        self.table.setHorizontalHeaderLabels(["工步序号", "工步状态", "小循环", "大循环"])
        self.table.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        # self.table.verticalHeader().setVisible(False)
        self.table.setRowCount(100)
                # 为表格视图安装事件过滤器
        self.table.viewport().installEventFilter(self)

        # 按钮布局
        self.clear_btn = QPushButton("清除数据")
        self.add_cycle_btn = QPushButton("添加循环")
        self.clear_all_cycles_btn = QPushButton("清除所有循环")
        self.export_btn = QPushButton("自动推断")

        btn_layout = QHBoxLayout()
        btn_layout.addWidget(self.clear_btn)
        btn_layout.addWidget(self.add_cycle_btn)
        btn_layout.addWidget(self.clear_all_cycles_btn)
        btn_layout.addWidget(self.export_btn)

        # 主布局
        main_layout = QVBoxLayout()
        main_layout.addWidget(self.table)
        main_layout.addLayout(btn_layout)

        central_widget = QWidget()
        central_widget.setLayout(main_layout)
        self.setCentralWidget(central_widget)

        # 连接信号与槽
        self.paste_action = QAction(self)
        self.paste_action.setShortcut(QKeySequence.Paste)
        self.paste_action.triggered.connect(self.paste_data)
        self.table.addAction(self.paste_action)

        self.add_cycle_btn.clicked.connect(self.add_cycle)
        self.clear_all_cycles_btn.clicked.connect(self.clear_all_cycles)
        self.clear_btn.clicked.connect(self.clear_data)
        self.export_btn.clicked.connect(self.infer_inner_cycle)

    def get_random_color(self):
        """从颜色池中随机选择一种颜色"""
        return random.choice(self.color_pool)

    def add_cycle(self):
        selected_rows = sorted(set(index.row() for index in self.table.selectedIndexes()))
        if not selected_rows:
            QMessageBox.warning(self, "警告", "请先选中要标记为循环的行！")
            return

        cycle_num = self.undo_stack.count() + 1  # 循环编号
        
        # 获取该循环的颜色，如果是新循环则随机分配
        if cycle_num not in self.cycle_colors:
            self.cycle_colors[cycle_num] = self.get_random_color()
        color = self.cycle_colors[cycle_num]

        # 保存操作前的标记状态和颜色状态
        prev_marks = {}
        prev_colors = {}
        for row in selected_rows:
            item = self.table.item(row, 2)
            prev_marks[row] = item.text() if item else ""
            prev_colors[row] = item.background() if item else QColor(255, 255, 255)

        # 创建并推入撤销命令
        cmd = MarkCycleCommand(self.table, selected_rows[0], selected_rows[-1], 
                              cycle_num, color, prev_marks, prev_colors)
        self.undo_stack.push(cmd)

    def clear_all_cycles(self):
        for row in range(self.table.rowCount()):
            self.table.setItem(row, 2, QTableWidgetItem(""))
        self.cycle_colors.clear()  # 清空循环颜色记录
        self.undo_stack.clear()

    def clear_data(self):
        self.show_borders = False
        self.table.clearContents()
        self.cycle_colors.clear()
        self.undo_stack.clear()

    def keyPressEvent(self, event):
        # Ctrl + Z 撤销
        if event.key() == Qt.Key_Z and (event.modifiers() & Qt.ControlModifier):
            self.undo_stack.undo()
        super().keyPressEvent(event)

    def paste_data(self):
        """从剪贴板粘贴数据到表格"""
        clipboard = QApplication.clipboard()
        mime_data = clipboard.mimeData()

        # 检查是否有文本数据
        if mime_data.hasText():
            # 获取剪贴板文本
            text = mime_data.text()
            # 将文本分割成行
            rows = text.split('\n')
            rows = [row for row in rows if row.strip()]  # 移除空行
            if not rows:
                QMessageBox.warning(self, "无数据", "剪贴板中未找到可粘贴的数据!")
                return

            # 计算行数和列数
            num_rows = len(rows)
            num_cols = max(len(row.split('\t')) for row in rows)

            # 设置表格大小
            self.table.setRowCount(num_rows)

            # 填充表格
            for row_idx, row in enumerate(rows):
                cols = row.split('\t')
                for col_idx, col_text in enumerate(cols[:2]):
                    item = QTableWidgetItem(col_text.strip())
                    self.table.setItem(row_idx, col_idx, item)

            self.statusBar().showMessage(f"成功粘贴 {num_rows} 行, {num_cols} 列数据 | 共 {num_rows * num_cols} 个单元格")
        else:
            QMessageBox.warning(self, "数据格式错误", "剪贴板中不包含文本数据，请复制 Excel 表格内容!")
        self.infer_outer_cycle()

    def infer_outer_cycle(self):
        """推测大循环"""
        rows = []
        for row in range(self.table.rowCount()):
            item = self.table.item(row, 0)
            if item:
                rows.append(item.text())
            else:
                break
        tagged_sequence = mark_repeating_sequences(rows)
        arr = np.array(tagged_sequence)
        # 比较相邻元素是否相等，找到变化点
        changes = arr[:-1] != arr[1:]
        change_points = np.where(changes)[0]
        # 生成区间端点并排序
        endpoints = sorted(set([0] + change_points.tolist() + [len(arr)-1]))
        # 生成连续区间
        output_rows = []
        before_point = None
        for i in range(1, len(endpoints)):
            start = endpoints[i-1] + 1 if before_point else endpoints[i-1]
            end = endpoints[i]
            before_point = (start, end, tagged_sequence[start])
            output_rows.append(before_point)
        
        self.special_row_ranges = output_rows
        self.show_borders = not self.show_borders
        self.table.viewport().update()


    def infer_inner_cycle(self):
        """推测小循环"""
        pos2cycle = {}
        one_count = 0
        value_positions = defaultdict(list)
        for row in range(self.table.rowCount()):
            item = self.table.item(row, 3)
            if item and item.text() == "1":
                source_item = self.table.item(row, 2)
                if source_item:
                    pos2cycle[one_count] = {"text": source_item.text(), "color": source_item.background()}
                one_count += 1
            if item and item.text():
                value_positions[item.text()].append(row)
        for row in range(self.table.rowCount()):
            item = self.table.item(row, 3)
            if item and item.text():
                if item.text() in ["0", "1"]:
                    continue
                else:
                    pos_index = value_positions[item.text()].index(row)
                    source_data = pos2cycle.get(pos_index)
                    if source_data:
                        text, color = source_data['text'], source_data['color']
                        target_item = QTableWidgetItem(text)
                        target_item.setBackground(color)
                        self.table.setItem(row, 2, target_item)

    def get_small_cycle_list(self) -> tuple[list, dict]:
        rows = []
        for row in range(self.table.rowCount()):
            cycle_item = self.table.item(row, 3)
            if cycle_item and cycle_item.text():
                big_cycle = int(cycle_item.text())
                small_cycle = self.table.item(row, 2)
                small_cycle_num = np.nan
                if small_cycle and small_cycle.text():
                    small_cycle_num = int(small_cycle.text()[3:])
                rows.append({"small_cycle": small_cycle_num, "big_cycle": big_cycle})
            else:
                rows.append({"small_cycle": np.nan, "big_cycle": np.nan})
        return rows


    def eventFilter(self, obj, event):
        # 监听表格视图的绘制事件
        if obj == self.table.viewport() and event.type() == QEvent.Paint:
            # 先让表格正常绘制
            result = super().eventFilter(obj, event)
            if self.show_borders:
                # 创建painter绘制自定义边框
                painter = QPainter(obj)
                if painter.isActive():
                    # 设置画笔：蓝色，2px粗，实线
                    pen = QPen(QColor("#1E90FF"), 2, Qt.SolidLine)
                    painter.setPen(pen)
                    # 绘制行范围的边框
                    for start_row, end_row, cycle_num in self.special_row_ranges:
                        # 获取起始行和结束行的单元格
                        start_item = self.table.item(start_row, 0)
                        end_item = self.table.item(end_row, 0)
                        if start_item and end_item:
                            # 获取起始单元格和结束单元格的矩形
                            start_rect = self.table.visualItemRect(start_item)
                            end_rect = self.table.visualItemRect(end_item)
                            # 计算包含整个行范围的矩形
                            combined_rect = start_rect.united(end_rect)
                            # 微调矩形，避免与原有边框重叠
                            combined_rect.adjust(1, 1, -1, -1)
                            # 绘制边框
                            painter.drawRect(combined_rect)
                        # 更新大循环圈数
                        for row in range(start_row, end_row+1):
                            item = QTableWidgetItem(str(cycle_num))
                            self.table.setItem(row, 3, item)
            return result
        # 其他事件交给默认处理
        return super().eventFilter(obj, event)


if __name__ == "__main__":
    app = QApplication([])
    window = WorkstepCycleTool()
    window.show()
    app.exec()
